#ifndef COMPLETE_BINARY_TREE_HPP
#define COMPLETE_BINARY_TREE_HPP

#include <vector>
#include <functional>
#include <stdexcept>
#include <cstddef>
#include <string>
#include <cassert>
#include <type_traits>
#include <memory>
#include <utility>
#include <algorithm>
#include <optional>

/**
 * @file CompleteBinaryTree.hpp
 * @brief Production-ready complete binary tree with heterogeneous node storage (C++17)
 * 
 * @details Features:
 * - Separate contiguous storage for leaf and internal nodes
 * - Strong exception safety guarantee for resize operations
 * - Allocator-aware design with C++17 allocator support
 * - STL-compatible interface with at() and operator[] accessors
 * - O(log n) path traversal with zero heap allocations
 */

/**
 * @brief Complete binary tree with separate storage for leaf and internal nodes
 * 
 * @tparam LeafType Type stored in leaf nodes. Must satisfy:
 *   - std::is_nothrow_destructible_v<LeafType>
 *   - std::is_nothrow_move_constructible_v<LeafType>
 * @tparam InternalType Type stored in internal nodes. Same requirements apply
 * @tparam Allocator Allocator type meeting C++17 Allocator requirements
 * 
 * @note Not thread-safe. External synchronization required.
 */
template<typename LeafType, typename InternalType, typename Allocator = std::allocator<void>>
class CompleteBinaryTree {
    // Type constraints
    static_assert(std::is_nothrow_destructible_v<LeafType>, 
                  "LeafType must be nothrow destructible");
    static_assert(std::is_nothrow_destructible_v<InternalType>, 
                  "InternalType must be nothrow destructible");
    static_assert(std::is_nothrow_move_constructible_v<LeafType>, 
                  "LeafType must be nothrow move constructible");
    static_assert(std::is_nothrow_move_constructible_v<InternalType>, 
                  "InternalType must be nothrow move constructible");

public:
    // Allocator support
    using allocator_type = Allocator;
    using leaf_allocator = typename std::allocator_traits<Allocator>::template rebind_alloc<LeafType>;
    using internal_allocator = typename std::allocator_traits<Allocator>::template rebind_alloc<InternalType>;

    /**
     * @brief Constructs an empty tree
     * @post leaf_count() == 0 && internal_count() == 0
     */
    explicit CompleteBinaryTree(const Allocator& alloc = Allocator()) noexcept
        : num_leaves_(0), num_internal_(0), leaves_(leaf_allocator(alloc)), 
          internals_(internal_allocator(alloc)) {}

    /**
     * @brief Constructs a tree with specified number of leaf nodes
     * @param num_leaves Number of leaf nodes. Must be > 0.
     * @param alloc Allocator instance
     * @throws std::invalid_argument If num_leaves == 0
     * @throws std::length_error If total node count exceeds SIZE_MAX
     */
    explicit CompleteBinaryTree(size_t num_leaves, const Allocator& alloc = Allocator());

    // Rule of five
    CompleteBinaryTree(const CompleteBinaryTree&) = default;
    CompleteBinaryTree& operator=(const CompleteBinaryTree&) = default;
    
    CompleteBinaryTree(CompleteBinaryTree&& other) noexcept(
        std::is_nothrow_move_constructible_v<LeafType> && 
        std::is_nothrow_move_constructible_v<InternalType> &&
        std::is_nothrow_move_constructible_v<Allocator>
    ) = default;
    
    CompleteBinaryTree& operator=(CompleteBinaryTree&& other) noexcept(
        std::is_nothrow_move_assignable_v<LeafType> && 
        std::is_nothrow_move_assignable_v<InternalType> &&
        std::is_nothrow_move_assignable_v<Allocator>
    ) = default;
    
    ~CompleteBinaryTree() = default;

    // Capacity Queries
    [[nodiscard]] size_t leaf_count() const noexcept { return num_leaves_; }
    [[nodiscard]] size_t internal_count() const noexcept { return num_internal_; }
    [[nodiscard]] size_t size() const noexcept { return num_leaves_ + num_internal_; }
    [[nodiscard]] bool empty() const noexcept { return num_leaves_ == 0; }
    [[nodiscard]] size_t depth() const noexcept { return calculateDepth(num_leaves_); }

    // Memory Management
    /** @brief Reserves storage for at least new_leaf_capacity leaves */
    void reserve(size_t new_leaf_capacity);
    
    /** @brief Reduces memory usage by freeing unused capacity */
    void shrink_to_fit() { leaves_.shrink_to_fit(); internals_.shrink_to_fit(); }

    // Element Access (STL-style)
    [[nodiscard]] LeafType& at(size_t leaf_index) { return getLeaf(leaf_index); }
    [[nodiscard]] const LeafType& at(size_t leaf_index) const { return getLeaf(leaf_index); }
    [[nodiscard]] InternalType& internal_at(size_t internal_index) { return getInternal(internal_index); }
    [[nodiscard]] const InternalType& internal_at(size_t internal_index) const { return getInternal(internal_index); }
    
    /** @brief Unchecked access. Undefined behavior if out of bounds. */
    [[nodiscard]] LeafType& operator[](size_t leaf_index) noexcept {
        assert(leaf_index < num_leaves_);
        return leaves_[leaf_index];
    }
    
    /** @brief Unchecked const access. Undefined behavior if out of bounds. */
    [[nodiscard]] const LeafType& operator[](size_t leaf_index) const noexcept {
        assert(leaf_index < num_leaves_);
        return leaves_[leaf_index];
    }

    /**
     * @brief Resizes the tree to new number of leaf nodes (strong exception-safe)
     * @param new_leaf_count New number of leaf nodes. Must be > 0.
     * @note Preserves existing elements following std::vector resize semantics
     */
    void resize(size_t new_leaf_count);

    /** @brief Clears all nodes */
    void clear() noexcept;

    /** @brief Swaps contents with another tree */
    void swap(CompleteBinaryTree& other) noexcept(
        std::allocator_traits<leaf_allocator>::propagate_on_container_swap::value ||
        std::allocator_traits<leaf_allocator>::is_always_equal::value
    );

    // Iterators
    using leaf_iterator = typename std::vector<LeafType, leaf_allocator>::iterator;
    using const_leaf_iterator = typename std::vector<LeafType, leaf_allocator>::const_iterator;
    using leaf_reverse_iterator = typename std::vector<LeafType, leaf_allocator>::reverse_iterator;
    using const_leaf_reverse_iterator = typename std::vector<LeafType, leaf_allocator>::const_reverse_iterator;

    leaf_iterator leaf_begin() noexcept { return leaves_.begin(); }
    leaf_iterator leaf_end() noexcept { return leaves_.end(); }
    const_leaf_iterator leaf_begin() const noexcept { return leaves_.begin(); }
    const_leaf_iterator leaf_end() const noexcept { return leaves_.end(); }
    const_leaf_iterator leaf_cbegin() const noexcept { return leaves_.cbegin(); }
    const_leaf_iterator leaf_cend() const noexcept { return leaves_.cend(); }
    leaf_reverse_iterator leaf_rbegin() noexcept { return leaves_.rbegin(); }
    leaf_reverse_iterator leaf_rend() noexcept { return leaves_.rend(); }
    const_leaf_reverse_iterator leaf_crbegin() const noexcept { return leaves_.crbegin(); }
    const_leaf_reverse_iterator leaf_crend() const noexcept { return leaves_.crend(); }

    using internal_iterator = typename std::vector<InternalType, internal_allocator>::iterator;
    using const_internal_iterator = typename std::vector<InternalType, internal_allocator>::const_iterator;
    using internal_reverse_iterator = typename std::vector<InternalType, internal_allocator>::reverse_iterator;
    using const_internal_reverse_iterator = typename std::vector<InternalType, internal_allocator>::const_reverse_iterator;

    internal_iterator internal_begin() noexcept { return internals_.begin(); }
    internal_iterator internal_end() noexcept { return internals_.end(); }
    const_internal_iterator internal_begin() const noexcept { return internals_.begin(); }
    const_internal_iterator internal_end() const noexcept { return internals_.end(); }
    const_internal_iterator internal_cbegin() const noexcept { return internals_.cbegin(); }
    const_internal_iterator internal_cend() const noexcept { return internals_.cend(); }
    internal_reverse_iterator internal_rbegin() noexcept { return internals_.rbegin(); }
    internal_reverse_iterator internal_rend() noexcept { return internals_.rend(); }
    const_internal_reverse_iterator internal_crbegin() const noexcept { return internals_.crbegin(); }
    const_internal_reverse_iterator internal_crend() const noexcept { return internals_.crend(); }

    // Path Operations
    [[nodiscard]] std::vector<std::reference_wrapper<InternalType>> path_to_root(size_t leaf_index);
    [[nodiscard]] std::vector<std::reference_wrapper<const InternalType>> path_to_root(size_t leaf_index) const;
    
    /** @brief Writes path nodes to output iterator without heap allocation */
    template<typename OutputIter>
    void path_to_root(size_t leaf_index, OutputIter out_iter);

    // Traversal Operations
    template<typename Func>
    void for_each_leaf(Func&& func) {
        static_assert(std::is_invocable_v<Func, LeafType&>, 
                      "Func must be invocable with LeafType&");
        std::for_each(leaves_.begin(), leaves_.end(), std::forward<Func>(func));
    }

    template<typename Func>
    void for_each_internal(Func&& func) {
        static_assert(std::is_invocable_v<Func, InternalType&>, 
                      "Func must be invocable with InternalType&");
        std::for_each(internals_.begin(), internals_.end(), std::forward<Func>(func));
    }

    /** @brief Applies function to each internal node on path to root (zero allocations) */
    template<typename Func>
    void for_each_on_path(size_t leaf_index, Func&& func);

    // Integrity Checking
    /** @brief Validates internal invariants */
    [[nodiscard]] bool validate_integrity() const noexcept;

private:
    // Index Conversions
    [[nodiscard]] inline size_t leaf_to_absolute(size_t leaf_index) const noexcept {
        return num_internal_ + leaf_index;
    }

    [[nodiscard]] inline std::optional<size_t> parent_of(size_t idx) const noexcept {
        return (idx == 0) ? std::nullopt : std::optional<size_t>{(idx - 1) / 2};
    }

    [[nodiscard]] inline bool is_leaf_index(size_t absolute_index) const noexcept {
        return absolute_index >= num_internal_;
    }

    // Validation and Access
    inline void validate_leaf_index(size_t leaf_index) const {
        if (leaf_index >= num_leaves_) {
            throw std::out_of_range(
                "Leaf index " + std::to_string(leaf_index) + 
                " out of range [0, " + std::to_string(num_leaves_ - 1) + "]"
            );
        }
    }

    inline void validate_internal_index(size_t internal_index) const {
        if (internal_index >= num_internal_) {
            throw std::out_of_range(
                "Internal index " + std::to_string(internal_index) + 
                " out of range [0, " + std::to_string(num_internal_ - 1) + "]"
            );
        }
    }

    [[nodiscard]] LeafType& getLeaf(size_t leaf_index) {
        validate_leaf_index(leaf_index);
        return leaves_[leaf_index];
    }

    [[nodiscard]] const LeafType& getLeaf(size_t leaf_index) const {
        validate_leaf_index(leaf_index);
        return leaves_[leaf_index];
    }

    [[nodiscard]] InternalType& getInternal(size_t internal_index) {
        validate_internal_index(internal_index);
        return internals_[internal_index];
    }

    [[nodiscard]] const InternalType& getInternal(size_t internal_index) const {
        validate_internal_index(internal_index);
        return internals_[internal_index];
    }

    // Integer-based depth calculation
    [[nodiscard]] static constexpr size_t calculateDepth(size_t num_leaves) noexcept {
        size_t depth = 0;
        for (size_t n = num_leaves; n > 0; n >>= 1) ++depth;
        return depth;
    }

    // Member variables
    size_t num_leaves_ = 0;
    size_t num_internal_ = 0;
    std::vector<LeafType, leaf_allocator> leaves_;
    std::vector<InternalType, internal_allocator> internals_;
};

// ==================== Implementation ====================

template<typename LeafType, typename InternalType, typename Allocator>
CompleteBinaryTree<LeafType, InternalType, Allocator>::CompleteBinaryTree(
    size_t num_leaves, const Allocator& alloc
) : leaves_(leaf_allocator(alloc)), internals_(internal_allocator(alloc)) {
    
    if (num_leaves == 0) {
        throw std::invalid_argument("Number of leaves must be positive");
    }

    // Overflow check: num_leaves + (num_leaves - 1) <= SIZE_MAX
    // This is equivalent to: 2 * num_leaves - 1 <= SIZE_MAX
    // Or: num_leaves <= (SIZE_MAX + 1) / 2
    // But to avoid overflow in SIZE_MAX + 1, we check: num_leaves > SIZE_MAX / 2 + 1
    if (num_leaves > SIZE_MAX / 2 + 1) {
        throw std::length_error("Total node count exceeds size_t range");
    }

    try {
        leaves_.resize(num_leaves);
        internals_.resize(num_leaves - 1);
    } catch (const std::bad_alloc& e) {
        throw std::runtime_error(
            std::string("Memory allocation failed for ") + 
            std::to_string(num_leaves) + " leaves: " + e.what()
        );
    }

    num_leaves_ = num_leaves;
    num_internal_ = num_leaves - 1;
}

template<typename LeafType, typename InternalType, typename Allocator>
void CompleteBinaryTree<LeafType, InternalType, Allocator>::reserve(size_t new_leaf_capacity) {
    if (new_leaf_capacity == 0) return;
    
    if (new_leaf_capacity > SIZE_MAX / 2 + 1) {
        throw std::length_error("Reserve capacity exceeds size_t range");
    }

    leaves_.reserve(new_leaf_capacity);
    internals_.reserve(new_leaf_capacity - 1);
}

template<typename LeafType, typename InternalType, typename Allocator>
void CompleteBinaryTree<LeafType, InternalType, Allocator>::resize(size_t new_leaf_count) {
    if (new_leaf_count == 0) {
        clear();
        return;
    }

    if (new_leaf_count == num_leaves_) return;

    if (new_leaf_count > SIZE_MAX / 2 + 1) {
        throw std::length_error("Total node count exceeds size_t range");
    }

    // Strong exception safety: prepare new storage first
    std::vector<LeafType, leaf_allocator> new_leaves(leaves_.get_allocator());
    std::vector<InternalType, internal_allocator> new_internals(internals_.get_allocator());

    new_leaves.resize(new_leaf_count);
    new_internals.resize(new_leaf_count - 1);

    // Move existing elements (follows std::vector::resize semantics)
    size_t copy_count = std::min(num_leaves_, new_leaf_count);
    if (copy_count > 0) {
        std::move(leaves_.begin(), leaves_.begin() + copy_count, new_leaves.begin());
        std::move(internals_.begin(), internals_.begin() + copy_count - 1, new_internals.begin());
    }

    // Commit changes (noexcept swap)
    leaves_.swap(new_leaves);
    internals_.swap(new_internals);
    num_leaves_ = new_leaf_count;
    num_internal_ = new_leaf_count - 1;
}

template<typename LeafType, typename InternalType, typename Allocator>
void CompleteBinaryTree<LeafType, InternalType, Allocator>::clear() noexcept {
    leaves_.clear();
    internals_.clear();
    num_leaves_ = 0;
    num_internal_ = 0;
}

template<typename LeafType, typename InternalType, typename Allocator>
void CompleteBinaryTree<LeafType, InternalType, Allocator>::swap(CompleteBinaryTree& other) noexcept(
    std::allocator_traits<leaf_allocator>::propagate_on_container_swap::value ||
    std::allocator_traits<leaf_allocator>::is_always_equal::value
) {
    using std::swap;
    swap(leaves_, other.leaves_);
    swap(internals_, other.internals_);
    swap(num_leaves_, other.num_leaves_);
    swap(num_internal_, other.num_internal_);
}

template<typename LeafType, typename InternalType, typename Allocator>
std::vector<std::reference_wrapper<InternalType>>
CompleteBinaryTree<LeafType, InternalType, Allocator>::path_to_root(size_t leaf_index) {
    std::vector<std::reference_wrapper<InternalType>> path;
    validate_leaf_index(leaf_index);
    path.reserve(depth());
    
    for (size_t idx = leaf_to_absolute(leaf_index); idx > 0; ) {
        auto parent = parent_of(idx);
        if (!parent) break;
        path.push_back(std::ref(internals_[*parent]));
        idx = *parent;
    }
    return path;
}

template<typename LeafType, typename InternalType, typename Allocator>
std::vector<std::reference_wrapper<const InternalType>>
CompleteBinaryTree<LeafType, InternalType, Allocator>::path_to_root(size_t leaf_index) const {
    std::vector<std::reference_wrapper<const InternalType>> path;
    validate_leaf_index(leaf_index);
    path.reserve(depth() - 1);
    
    for (size_t idx = leaf_to_absolute(leaf_index); idx > 0; ) {
        auto parent = parent_of(idx);
        if (!parent) break;
        path.push_back(std::cref(internals_[*parent]));
        idx = *parent;
    }
    return path;
}

template<typename LeafType, typename InternalType, typename Allocator>
template<typename OutputIter>
void CompleteBinaryTree<LeafType, InternalType, Allocator>::path_to_root(
    size_t leaf_index, OutputIter out_iter) {
    
    validate_leaf_index(leaf_index);
    for (size_t idx = leaf_to_absolute(leaf_index); idx > 0; ) {
        auto parent = parent_of(idx);
        if (!parent) break;
        *out_iter++ = std::ref(internals_[*parent]);
        idx = *parent;
    }
}

template<typename LeafType, typename InternalType, typename Allocator>
template<typename Func>
void CompleteBinaryTree<LeafType, InternalType, Allocator>::for_each_on_path(
    size_t leaf_index, Func&& func) {
    
    static_assert(std::is_invocable_v<Func, InternalType&>, 
                  "Func must be invocable with InternalType&");
    
    validate_leaf_index(leaf_index);
    if (num_internal_ == 0) return;

    for (size_t idx = leaf_to_absolute(leaf_index); idx > 0; ) {
        auto parent = parent_of(idx);
        if (!parent) break;
        assert(*parent < num_internal_);
        std::forward<Func>(func)(internals_[*parent]);
        idx = *parent;
    }
}

template<typename LeafType, typename InternalType, typename Allocator>
bool CompleteBinaryTree<LeafType, InternalType, Allocator>::validate_integrity() const noexcept {
    if (num_leaves_ == 0) {
        return num_internal_ == 0 && leaves_.empty() && internals_.empty();
    }
    
    if (num_internal_ != num_leaves_ - 1) return false;
    if (leaves_.size() != num_leaves_) return false;
    if (internals_.size() != num_internal_) return false;
    if (num_leaves_ > SIZE_MAX / 2 + 1) return false;
    if (depth() != calculateDepth(num_leaves_)) return false;
    
    // Allocator consistency check
    return leaves_.get_allocator() == internals_.get_allocator();
}

template<typename LeafType, typename InternalType, typename Allocator>
void swap(CompleteBinaryTree<LeafType, InternalType, Allocator>& lhs, 
          CompleteBinaryTree<LeafType, InternalType, Allocator>& rhs) 
    noexcept(noexcept(lhs.swap(rhs))) {
    lhs.swap(rhs);
}

#endif // COMPLETE_BINARY_TREE_HPP